import cv2
import numpy as np
import json
import os
import tkinter as tk
from tkinter import messagebox, filedialog
import platform

class RemapCalibrationTool:
    """
    Interactive tool for calibrating perspective transform parameters for remapping.
    Allows dragging control points and saving/loading parameters to/from JSON.
    Supports vertical stretching/compression and horizontal compression with black fill.
    """
    
    def __init__(self, 
                 camera_index=0, 
                 width=640, 
                 height=480, 
                 fps=30,
                 json_path="remap_params.json"):
        """
        Initialize the calibration tool.
        
        Args:
            camera_index: Camera device index (default: 0)
            width: Camera width resolution (default: 640)
            height: Camera height resolution (default: 480)
            fps: Camera frames per second (default: 30)
            json_path: Path to save/load JSON parameters (default: "remap_params.json")
        """
        self.camera_index = camera_index
        self.width = width
        self.height = height
        self.fps = fps
        self.json_path = json_path
        
        # Initialize camera
        system = platform.system()
        if system == "Windows":
            self.cap = cv2.VideoCapture(self.camera_index, cv2.CAP_DSHOW)
        elif system == "Linux":
            self.cap = cv2.VideoCapture(self.camera_index, cv2.CAP_V4L2)
        else:
            self.cap = cv2.VideoCapture(self.camera_index)
            
        # Set camera properties
        self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
        self.cap.set(cv2.CAP_PROP_FPS, self.fps)
        
        # Check if camera opened successfully
        if not self.cap.isOpened():
            raise ValueError(f"Failed to open camera at index {camera_index}")
            
        # Get actual camera dimensions (may be different from requested)
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # Initialize parameters with default values
        self.params = {
            'SRC_TOP_LEFT_X': 0.2,
            'SRC_TOP_LEFT_Y': 0.2,
            'SRC_TOP_RIGHT_X': 0.8,
            'SRC_TOP_RIGHT_Y': 0.2,
            'SRC_BOTTOM_RIGHT_X': 0.8,
            'SRC_BOTTOM_RIGHT_Y': 0.8,
            'SRC_BOTTOM_LEFT_X': 0.2,
            'SRC_BOTTOM_LEFT_Y': 0.8,
            'DST_RIGHT_MARGIN': 0,
            'VERTICAL_STRETCH': 1.0,  # Default to no stretching (1.0x)
            'HORIZONTAL_COMPRESSION': 1.0  # Default to no compression (1.0x)
        }
        
        # Try to load parameters from JSON file if it exists
        self.load_params()
        
        # Initialize control points
        self.update_control_points()
        
        # State variables for point dragging
        self.dragging = False
        self.selected_point = -1
        self.drag_threshold = 20  # Distance in pixels to detect point selection
        
        # Create window for calibration
        cv2.namedWindow('Calibration', cv2.WINDOW_NORMAL)
        cv2.setMouseCallback('Calibration', self.on_mouse)
        
        # Create window for preview
        cv2.namedWindow('Remapped', cv2.WINDOW_NORMAL)
        
        # Initialize instructions overlay
        self.update_instructions()
        
        # Running flag
        self.running = True
        
    def update_instructions(self):
        """Update the instructions overlay with current parameter values."""
        self.instructions = [
            "Drag the RED POINTS to adjust the perspective transform",
            f"Press + or - to adjust vertical stretch (current: {self.params['VERTICAL_STRETCH']:.1f}x)",
            f"Press [ or ] to adjust horizontal compression (current: {self.params['HORIZONTAL_COMPRESSION']:.1f}x)",
            "Press S to save parameters to JSON",
            "Press L to load parameters from JSON",
            "Press R to reset parameters to default",
            "Press Q or ESC to quit"
        ]
        
    def update_control_points(self):
        """Update control points based on current parameters."""
        self.control_points = np.array([
            [int(self.width * self.params['SRC_TOP_LEFT_X']), 
             int(self.height * self.params['SRC_TOP_LEFT_Y'])],
            [int(self.width * self.params['SRC_TOP_RIGHT_X']), 
             int(self.height * self.params['SRC_TOP_RIGHT_Y'])],
            [int(self.width * self.params['SRC_BOTTOM_RIGHT_X']), 
             int(self.height * self.params['SRC_BOTTOM_RIGHT_Y'])],
            [int(self.width * self.params['SRC_BOTTOM_LEFT_X']), 
             int(self.height * self.params['SRC_BOTTOM_LEFT_Y'])]
        ], dtype=np.float32)
        
        # Calculate destination points
        self.dest_points = np.array([
            [0, 0],
            [self.width - self.params['DST_RIGHT_MARGIN'], 0],
            [self.width - self.params['DST_RIGHT_MARGIN'], self.height],
            [0, self.height]
        ], dtype=np.float32)
        
        # Calculate perspective transform matrix
        self.M = cv2.getPerspectiveTransform(self.control_points, self.dest_points)
    
    def on_mouse(self, event, x, y, flags, param):
        """Handle mouse events for dragging control points."""
        if event == cv2.EVENT_LBUTTONDOWN:
            # Check if we clicked near a control point
            for i, point in enumerate(self.control_points):
                dist = np.sqrt((point[0] - x)**2 + (point[1] - y)**2)
                if dist < self.drag_threshold:
                    self.dragging = True
                    self.selected_point = i
                    break
        
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.dragging and self.selected_point >= 0:
                # Update the selected control point position
                self.control_points[self.selected_point] = [x, y]
                
                # Update parameters based on new control points
                point_names = ['SRC_TOP_LEFT', 'SRC_TOP_RIGHT', 
                               'SRC_BOTTOM_RIGHT', 'SRC_BOTTOM_LEFT']
                self.params[f'{point_names[self.selected_point]}_X'] = x / self.width
                self.params[f'{point_names[self.selected_point]}_Y'] = y / self.height
                
                # Recalculate transform matrix
                self.M = cv2.getPerspectiveTransform(self.control_points, self.dest_points)
        
        elif event == cv2.EVENT_LBUTTONUP:
            self.dragging = False
            self.selected_point = -1
    
    def draw_control_points(self, frame):
        """Draw control points and connecting lines on the frame."""
        frame_with_points = frame.copy()
        
        # Draw lines connecting the points
        cv2.line(frame_with_points, 
                tuple(self.control_points[0].astype(int)), 
                tuple(self.control_points[1].astype(int)), 
                (0, 255, 0), 2)
        cv2.line(frame_with_points, 
                tuple(self.control_points[1].astype(int)), 
                tuple(self.control_points[2].astype(int)), 
                (0, 255, 0), 2)
        cv2.line(frame_with_points, 
                tuple(self.control_points[2].astype(int)), 
                tuple(self.control_points[3].astype(int)), 
                (0, 255, 0), 2)
        cv2.line(frame_with_points, 
                tuple(self.control_points[3].astype(int)), 
                tuple(self.control_points[0].astype(int)), 
                (0, 255, 0), 2)
        
        # Draw circles at the control points
        for i, point in enumerate(self.control_points):
            center = tuple(point.astype(int))
            # Make the selected point bigger
            radius = 8 if self.selected_point == i else 5
            # Use a different color for the selected point
            color = (0, 0, 255) if self.selected_point == i else (0, 0, 255)
            cv2.circle(frame_with_points, center, radius, color, -1)
            
            # Add point labels
            label = ["Top Left", "Top Right", "Bottom Right", "Bottom Left"][i]
            cv2.putText(frame_with_points, label, 
                       (center[0] + 10, center[1] + 10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
        
        # Add instructions overlay
        y_pos = 30
        for instruction in self.instructions:
            cv2.putText(frame_with_points, instruction, 
                       (10, y_pos), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
            y_pos += 25
        
        return frame_with_points
    
    def apply_remap(self, frame):
        """Apply perspective transform to remap the frame, with optional stretching/compression."""
        # First apply the perspective transform
        remapped = cv2.warpPerspective(frame, self.M, (self.width, self.height))
        
        # Apply vertical stretching/compression if needed
        if self.params['VERTICAL_STRETCH'] != 1.0:
            remapped = self.stretch_image_vertically(remapped, self.params['VERTICAL_STRETCH'])
        
        # Apply horizontal compression if needed
        if self.params['HORIZONTAL_COMPRESSION'] != 1.0:
            remapped = self.compress_image_horizontally(remapped, self.params['HORIZONTAL_COMPRESSION'])
        
        return remapped
    
    def stretch_image_vertically(self, image, factor):
        """
        Stretch or compress an image vertically by the given factor.
        
        Args:
            image: The input image
            factor: The stretching factor (>1.0 stretches, <1.0 compresses)
            
        Returns:
            The vertically stretched/compressed image
        """
        h, w = image.shape[:2]
        
        # Calculate new height
        new_h = int(h * factor)
        
        # Resize the image with the new height
        stretched = cv2.resize(image, (w, new_h))
        
        # If the stretching makes the image larger than the original height,
        # crop it to the original height (centered)
        if new_h > h:
            start_y = (new_h - h) // 2
            stretched = stretched[start_y:start_y+h, :]
        # If the stretching makes the image smaller, pad it with black
        elif new_h < h:
            pad_top = (h - new_h) // 2
            pad_bottom = h - new_h - pad_top
            stretched = cv2.copyMakeBorder(stretched, pad_top, pad_bottom, 0, 0, 
                                          cv2.BORDER_CONSTANT, value=(0, 0, 0))
        
        return stretched
    
    def compress_image_horizontally(self, image, factor):
        """
        Compress an image horizontally by the given factor and fill with black.
        
        Args:
            image: The input image
            factor: The compression factor (<1.0 compresses, >1.0 would stretch)
            
        Returns:
            The horizontally compressed image with black fill
        """
        h, w = image.shape[:2]
        
        # Calculate new width (compressed)
        new_w = int(w * factor)
        
        # Resize the image to the new width (maintaining height)
        compressed = cv2.resize(image, (new_w, h))
        
        # Create a black canvas of original size
        result = np.zeros((h, w, 3), dtype=np.uint8)
        
        # Calculate padding to center the compressed image
        pad_left = (w - new_w) // 2
        
        # Place the compressed image in the center
        result[:, pad_left:pad_left+new_w] = compressed
        
        return result
    
    def save_params(self, custom_path=None):
        """Save parameters to JSON file."""
        try:
            # Use custom path if provided, otherwise use default
            save_path = custom_path if custom_path else self.json_path
            
            # Format parameters with float precision
            formatted_params = {k: round(float(v), 6) for k, v in self.params.items()}
            
            with open(save_path, 'w') as f:
                json.dump(formatted_params, f, indent=4)
            
            print(f"Parameters saved to {save_path}")
            return True
        except Exception as e:
            print(f"Error saving parameters: {e}")
            return False
    
    def load_params(self, custom_path=None):
        """Load parameters from JSON file."""
        try:
            # Use custom path if provided, otherwise use default
            load_path = custom_path if custom_path else self.json_path
            
            # Check if file exists
            if not os.path.exists(load_path):
                print(f"File {load_path} does not exist. Using default parameters.")
                return False
            
            with open(load_path, 'r') as f:
                loaded_params = json.load(f)
            
            # Update parameters
            self.params.update(loaded_params)
            
            # Update control points based on loaded parameters
            self.update_control_points()
            
            # Update instructions with loaded values
            self.update_instructions()
            
            print(f"Parameters loaded from {load_path}")
            return True
        except Exception as e:
            print(f"Error loading parameters: {e}")
            return False
    
    def reset_params(self):
        """Reset parameters to default values."""
        self.params = {
            'SRC_TOP_LEFT_X': 0.2,
            'SRC_TOP_LEFT_Y': 0.2,
            'SRC_TOP_RIGHT_X': 0.8,
            'SRC_TOP_RIGHT_Y': 0.2,
            'SRC_BOTTOM_RIGHT_X': 0.8,
            'SRC_BOTTOM_RIGHT_Y': 0.8,
            'SRC_BOTTOM_LEFT_X': 0.2,
            'SRC_BOTTOM_LEFT_Y': 0.8,
            'DST_RIGHT_MARGIN': 0,
            'VERTICAL_STRETCH': 1.0,
            'HORIZONTAL_COMPRESSION': 1.0
        }
        self.update_control_points()
        self.update_instructions()
        print("Parameters reset to default values")
    
    def show_save_dialog(self):
        """Show a file dialog to select where to save parameters."""
        # Create a temporary Tkinter root
        root = tk.Tk()
        root.withdraw()  # Hide the root window
        
        # Show save file dialog
        file_path = filedialog.asksaveasfilename(
            defaultextension=".json",
            filetypes=[("JSON files", "*.json"), ("All files", "*.*")],
            initialfile=self.json_path
        )
        
        # Destroy the root
        root.destroy()
        
        if file_path:
            return self.save_params(file_path)
        return False
    
    def show_load_dialog(self):
        """Show a file dialog to select where to load parameters from."""
        # Create a temporary Tkinter root
        root = tk.Tk()
        root.withdraw()  # Hide the root window
        
        # Show open file dialog
        file_path = filedialog.askopenfilename(
            filetypes=[("JSON files", "*.json"), ("All files", "*.*")],
            initialfile=self.json_path
        )
        
        # Destroy the root
        root.destroy()
        
        if file_path:
            return self.load_params(file_path)
        return False
    
    def update_vertical_stretch(self, delta):
        """Update the vertical stretch factor by the given delta."""
        # Update the stretch factor
        new_stretch = self.params['VERTICAL_STRETCH'] + delta
        
        # Apply reasonable limits (e.g., between 0.5x and 3.0x)
        new_stretch = max(0.5, min(3.0, new_stretch))
        
        # Update the parameter
        self.params['VERTICAL_STRETCH'] = new_stretch
        
        # Update the instructions
        self.update_instructions()
        
        print(f"Vertical stretch set to {new_stretch:.1f}x")
    
    def update_horizontal_compression(self, delta):
        """Update the horizontal compression factor by the given delta."""
        # Update the compression factor
        new_compression = self.params['HORIZONTAL_COMPRESSION'] + delta
        
        # Apply reasonable limits (e.g., between 0.2x and 1.5x)
        new_compression = max(0.2, min(1.5, new_compression))
        
        # Update the parameter
        self.params['HORIZONTAL_COMPRESSION'] = new_compression
        
        # Update the instructions
        self.update_instructions()
        
        print(f"Horizontal compression set to {new_compression:.1f}x")
    
    def run(self):
        """Run the calibration tool main loop."""
        while self.running and self.cap.isOpened():
            # Capture frame
            ret, frame = self.cap.read()
            if not ret:
                print("Failed to capture frame")
                break
            
            # Draw control points on the frame
            frame_with_points = self.draw_control_points(frame)
            
            # Apply perspective transform and stretching/compression
            remapped_frame = self.apply_remap(frame)
            
            # Display frames
            cv2.imshow('Calibration', frame_with_points)
            cv2.imshow('Remapped', remapped_frame)
            
            # Handle keyboard input
            key = cv2.waitKey(1) & 0xFF
            
            if key == ord('q') or key == 27:  # 'q' or ESC
                self.running = False
            elif key == ord('s'):  # 's' to save
                self.show_save_dialog()
            elif key == ord('l'):  # 'l' to load
                self.show_load_dialog()
            elif key == ord('r'):  # 'r' to reset
                self.reset_params()
            elif key == ord('+') or key == ord('='):  # '+' to increase vertical stretch
                self.update_vertical_stretch(0.1)
            elif key == ord('-'):  # '-' to decrease vertical stretch
                self.update_vertical_stretch(-0.1)
            elif key == ord(']'):  # ']' to decrease horizontal compression
                self.update_horizontal_compression(-0.1)
            elif key == ord('['):  # '[' to increase horizontal compression
                self.update_horizontal_compression(0.1)
            elif key == ord('}'):  # '}' to decrease horizontal compression more
                self.update_horizontal_compression(-0.5)
            elif key == ord('{'):  # '{' to increase horizontal compression more
                self.update_horizontal_compression(0.5)
                
        # Clean up
        self.cap.release()
        cv2.destroyAllWindows()
    
    def get_params(self):
        """Get the current parameters."""
        return self.params
    
    def __del__(self):
        """Clean up resources when the object is destroyed."""
        if hasattr(self, 'cap') and self.cap.isOpened():
            self.cap.release()
        cv2.destroyAllWindows()


def main():
    """Main function to run the calibration tool."""
    try:
        # Create and run the calibration tool
        calibrator = RemapCalibrationTool(
            camera_index=1,
            width=640,
            height=480,
            fps=30,
            json_path="remap_params.json"
        )
        
        # Run the calibration loop
        calibrator.run()
        
    except Exception as e:
        print(f"Error: {e}")
        
    finally:
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()